from numpy import *
from KFacce import axisFilter
from KFacce import axisFilterBias
from LoadTUMData import *
from AvgAcce import *
from FilterVandP import *

# def main():


def dataInitialize(getGroup):

    q0_init, q1_init, q2_init, q3_init = getGroup[6][0], getGroup[6][1], getGroup[6][2], getGroup[6][3]
    Pinit_x, Pinit_y, Pinit_z = getGroup[3][0], getGroup[3][1], getGroup[3][2]
    Vinit_x, Vinit_y, Vinit_z = (Pinit_x - getGroup[4][0]) / 0.008, (Pinit_y - getGroup[4][1]) / 0.008, (Pinit_z - getGroup[4][2]) / 0.008
    Plast_x, Plast_y, Plast_z = getGroup[4][0], getGroup[4][1], getGroup[4][2]

    initialize = [q0_init, q1_init, q2_init, q3_init, Pinit_x, Pinit_y, Pinit_z, Vinit_x, Vinit_y, Vinit_z,
                  Plast_x, Plast_y, Plast_z]
    return initialize

def KFfilterBias(getGroup):
    Tum_accex, Tum_accey, Tum_accez, Tum_angux, Tum_anguy, Tum_anguz = [], [], [], [], [], []
    # time, Tum_accex, Tum_accey, Tum_accez, Tum_angux, Tum_anguy, Tum_anguz = loadTumImu()
    Tum_accex, Tum_accey, Tum_accez, Tum_angux, Tum_anguy, Tum_anguz = getGroup[7], getGroup[8], getGroup[9], \
                                                                       getGroup[10], getGroup[11], getGroup[12]


    initialize = dataInitialize(getGroup)
    print(initialize)   ###

    q0_init, q1_init, q2_init, q3_init = initialize[0], initialize[1], initialize[2], initialize[3]

    ### a028_3 initial values
    Q_before = array([[q0_init], [q1_init], [q2_init], [q3_init]])

    datacnt = 0

    delta_t = 0.02
    Z = array([0, 0, 0])

    K = array([[1, 0, 0],
               [0, 1, 0],
               [0, 0, 1],
               [0, 0, 0],
               [0, 0, 0],
               [0, 0, 0]])
    P = array([[1, 0, 0, 0, 0, 0],
               [0, 1, 0, 0, 0, 0],
               [0, 0, 1, 0, 0, 0],
               [0, 0, 0, 1, 0, 0],
               [0, 0, 0, 0, 1, 0],
               [0, 0, 0, 0, 0, 1]])
    PP = array([[1, 0, 0, 0, 0, 0],
                [0, 1, 0, 0, 0, 0],
                [0, 0, 1, 0, 0, 0],
                [0, 0, 0, 1, 0, 0],
                [0, 0, 0, 0, 1, 0],
                [0, 0, 0, 0, 0, 1]])

    A = array([[1, 0, 0, delta_t, 0, 0],
               [0, 1, 0, 0, delta_t, 0],
               [0, 0, 1, 0, 0, delta_t],
               [0, 0, 0, 1, 0, 0],
               [0, 0, 0, 0, 1, 0],
               [0, 0, 0, 0, 0, 1]])
    H = array([[1, 0, 0, 0, 0, 0],
               [0, 1, 0, 0, 0, 0],
               [0, 0, 1, 0, 0, 0]])
    F = array([[0.5 * delta_t * delta_t, 0, 0],
               [0, 0.5 * delta_t * delta_t, 0],
               [0, 0, 0.5 * delta_t * delta_t],
               [delta_t, 0, 0],
               [0, delta_t, 0],
               [0, 0, delta_t]])

    Q = 0.001 * identity(3)
    R = 0.01 * identity(3)

    w1 = Tum_angux[datacnt]
    w2 = Tum_anguy[datacnt]
    w3 = Tum_anguz[datacnt]
    X = array([[w1], [w2], [w3], [0], [0], [0]])

    fileOut = open('Android1st.csv', 'a')
    csv_writer = csv.writer(fileOut)
    csv_writer.writerow(
        ['wp', 'wq', 'wr', 'q0', 'q1', 'q2', 'q3', 'a0', 'a1', 'a2', 'raw_a0', 'raw_a1', 'raw_a2', 'Px', 'Py', 'Pz',
         'Raw_Px', 'Raw_Py', 'Raw_Pz'])
    # fileMeta = open('metaData.csv', 'a')

## jiashang bias axis acce guolv
##
    acce_bias = 0
    Q_a = 1
    Q_a_bias = 1
    p_a = 0.5
    p_a_bias = 0.5

    acce_state = array([[0], [acce_bias]])

    F_acce = array([[1, 0], [-1, 0]])
    B_acce = array([0])

    Q_acce = array([[Q_a, 0], [0, Q_a_bias]])

    H_acce = array([1, 0])
    R_acce = array([[0.01]])

    P_acce = array([[p_a, 0], [0, p_a_bias]])
    PP_acce = array([[], []])
    K_acce = array([[0.5], [0.5]])

    # AxisFilterBias(acce_state, acce_measure, F_acce, Q_acce, H_acce, R_acce, P_acce, PP_acce, K_acce)

    paramAcce_x = [acce_state, 0, F_acce, Q_acce, H_acce, R_acce, P_acce, K_acce]
    paramAcce_y = [acce_state, 0, F_acce, Q_acce, H_acce, R_acce, P_acce, K_acce]
    paramAcce_z = [acce_state, 0, F_acce, Q_acce, H_acce, R_acce, P_acce, K_acce]
##
##

    # K_acce = array([[1]])
    # P_acce = array([[1]])
    # PP_acce = array([[1]])

    # A_acce = array([[1]])
    # H_acce = array([[1]])
    # F_acce = array([[delta_t]])
    # Q_acce = array([[100]])
    # R_acce = array([[0.01]])

    Pinit_x, Pinit_y, Pinit_z = initialize[4], initialize[5], initialize[6]

    #### Error: Tum_accey[0] and Tum_accez[0]
    #### PreviousError: Tum_acce did not start from 8500, but it seems Okay
    # paramAcce_x = [0, 0, K_acce, P_acce, PP_acce, A_acce, H_acce, F_acce, array([[0.005]]), R_acce]
    # paramAcce_y = [0, 0, K_acce, P_acce, PP_acce, A_acce, H_acce, F_acce, array([[0.1]]), R_acce]
    # paramAcce_z = [0, 0, K_acce, P_acce, PP_acce, A_acce, H_acce, F_acce, Q_acce, R_acce]


    #### delta_t here is wrong in GT_data
    Vinit_x = initialize[7]
    Vinit_y = initialize[8]
    Vinit_z = initialize[9]

    ###try GT acce input


    RawP_x = Pinit_x
    RawP_y = Pinit_y
    RawP_z = Pinit_z

    RawV_x = Vinit_x
    RawV_y = Vinit_y
    RawV_z = Vinit_z

    ### aver3 aver5 & aver10
    AverX, AverY, AverZ = 0, 0, 0
    Q3x = queue.Queue(3)
    Q3y = queue.Queue(3)
    Q3z = queue.Queue(3)
    Q5x = queue.Queue(5)
    Q5y = queue.Queue(5)
    Q5z = queue.Queue(5)
    Q10x = queue.Queue(10)
    Q10y = queue.Queue(10)
    Q10z = queue.Queue(10)

    LastV_x, LastV_y, LastV_z = (initialize[7]), initialize[8], initialize[9]
    LastP_x, LastP_y, LastP_z = initialize[10], initialize[11], initialize[12]


    drift_x = 0
    drift_y = 0

    while (datacnt < len(Tum_accex) - 1 and datacnt < len(Tum_angux) - 1):


        angu = array([w1, w2, w3])

        # X = array([w1, w2, w3, 0, 0, 0])

        Z = array([[w1], [w2], [w3]])
        # print(shape(Z))

        ###Z represents for measurement

        X = dot(A, X) + dot(K, (Z - dot(H, dot(A, X))))

        # print(shape(X))

        K = dot(PP, dot(H.T, linalg.inv(dot(H, dot(PP, H.T)) + R)))
        PP = dot(A, dot(P, A.T)) + dot(F, dot(Q, F.T))
        P = PP - dot(K, dot(H, PP))

        wx = X[0][0]
        wy = X[1][0]
        wz = X[2][0]

        datacnt += 1
        w1 = Tum_angux[datacnt]
        w2 = Tum_anguy[datacnt]
        w3 = Tum_anguz[datacnt]

        # print((wx, wy, wz))
        # print((w1, w2, w3))

        anguToQ = [[0, -wx, -wy, -wz],
                   [wx, 0, wz, -wy],
                   [wy, -wz, 0, wx],
                   [wz, wy, -wx, 0]]
        qt = 0.5 * delta_t * dot(anguToQ, Q_before) + Q_before
        # qt[0][0] =

        ### insert acce filter this place

        q0, q1, q2, q3 = qt[0][0], qt[1][0], qt[2][0], qt[3][0]
        Tbi = array([[1 - 2 * (q2 ** 2 + q3 ** 2), 2 * (q1 * q2 - q0 * q3), 2 * (q1 * q3 + q0 * q2)],
                     [2 * (q1 * q2 + q0 * q3), 1 - 2 * (q1 ** 2 + q3 ** 2), 2 * (q2 * q3 - q0 * q1)],
                     [2 * (q1 * q3 - q0 * q2), 2 * (q2 * q3 + q0 * q1), 1 - 2 * (q1 ** 2 + q2 ** 2)]])

        ### using ground truth Q for rotation matrix
        # Gt_Tbi = array([[1 - 2 * (Gt_q2 ** 2 + Gt_q3 ** 2), 2 * (Gt_q1 * Gt_q2 - Gt_q0 * Gt_q3), 2 * (Gt_q1 * Gt_q3 + Gt_q0 * Gt_q2)],
        #                 [2 * (Gt_q1 * Gt_q2 + Gt_q0 * Gt_q3), 1 - 2 * (Gt_q1 ** 2 + Gt_q3 ** 2), 2 * (Gt_q2 * Gt_q3 - Gt_q0 * Gt_q1)],
        #                 [2 * (Gt_q1 * Gt_q3 - Gt_q0 * Gt_q2), 2 * (Gt_q2 * Gt_q3 + Gt_q0 * Gt_q1), 1 - 2 * (Gt_q1 ** 2 + Gt_q2 ** 2)]])

        acceX, acceY, acceZ, AverX, AverY, AverZ = filterNum(10, Tum_accex[datacnt], Tum_accey[datacnt], Tum_accez[datacnt],
                                                             Q10x, Q10y, Q10z, AverX, AverY, AverZ)
        acce = array([acceX, acceY, acceZ])
        # acce = array([Tum_accex[datacnt], Tum_accey[datacnt], Tum_accez[datacnt]])

        # if datacnt == 8501:
        #     acce = array([Ainit_x, Ainit_y, Ainit_z]) juede buyong chushihua zhege Acce

        # acce = array([Tango_a0[datacnt], Tango_a1[datacnt], Tango_a2[datacnt]])
        # acce = dot(Gt_Tbi, acce)
        print(qt)
        acce = dot(Tbi, acce)
        print(acce)
        acce[2] = acce[2] - 9.81


###### tianjia Drift buchang zheli
        acce[0] = acce[0] + drift_x
        acce[1] = acce[1] + drift_y



        # Gt_q0, Gt_q1, Gt_q2, Gt_q3 = tango_ori[datacnt][0], tango_ori[datacnt][1], tango_ori[datacnt][2], tango_ori[datacnt][3]

        rawAcce_x = acce[0]
        rawAcce_y = acce[1]
        rawAcce_z = acce[2]

        RawV_x = RawV_x + rawAcce_x * delta_t
        RawV_y = RawV_y + rawAcce_y * delta_t
        RawV_z = RawV_z + rawAcce_z * delta_t

        RawP_x = RawP_x + RawV_x * delta_t + 0.5 * rawAcce_x * delta_t * delta_t
        RawP_y = RawP_y + RawV_y * delta_t + 0.5 * rawAcce_y * delta_t * delta_t
        RawP_z = RawP_z + RawV_z * delta_t + 0.5 * rawAcce_z * delta_t * delta_t


        #### yao tiaozheng guolv jizhi yinggai congzheli fenkai

        # ## Integrate to V and then filter
        # Vinit_x = Vinit_x + acce[0] * delta_t
        # Vinit_y = Vinit_y + acce[1] * delta_t
        # Vinit_z = Vinit_z + (acce[2]) * delta_t

        # paramAcce_x[1] = RawV_x
        # paramAcce_y[1] = RawV_y
        # paramAcce_z[1] = RawV_z

        # ## Integrate to P and then filter
        # Pinit_x = Pinit_x + Vinit_x * delta_t + 0.5 * acce[0] * delta_t * delta_t
        # Pinit_y = Pinit_y + Vinit_y * delta_t + 0.5 * acce[1] * delta_t * delta_t
        # Pinit_z = Pinit_z + Vinit_z * delta_t + 0.5 * (acce[2]) * delta_t * delta_t

        paramAcce_x[1] = RawP_x
        paramAcce_y[1] = RawP_y
        paramAcce_z[1] = RawP_z

        # paramAcce_x[1] = acce[0]
        # paramAcce_y[1] = acce[1]
        # paramAcce_z[1] = acce[2]

        paramAcce_x = axisFilterBias(paramAcce_x[0], paramAcce_x[1], paramAcce_x[2], paramAcce_x[3], paramAcce_x[4],
                                 paramAcce_x[5], paramAcce_x[6], paramAcce_x[7])
        paramAcce_y = axisFilterBias(paramAcce_y[0], paramAcce_y[1], paramAcce_y[2], paramAcce_y[3], paramAcce_y[4],
                                 paramAcce_y[5], paramAcce_y[6], paramAcce_x[7])
        paramAcce_z = axisFilterBias(paramAcce_z[0], paramAcce_z[1], paramAcce_z[2], paramAcce_z[3], paramAcce_z[4],
                                 paramAcce_z[5], paramAcce_z[6], paramAcce_z[7])

        #     acceX, acceY, acceZ = filterNum(5, paramAcce_x[0][0][0], paramAcce_y[0][0][0], paramAcce_z[0][0][0], AverX, AverY, AverZ)
        # ### integrate after 3 points filter, smooth/average acces after KF processed
        #     Vinit_x = Vinit_x + acceX * delta_t
        #     Vinit_y = Vinit_y + acceY * delta_t
        #     Vinit_z = Vinit_z + (acceZ) * delta_t
        #
        #     Pinit_x = Pinit_x + Vinit_x * delta_t + 0.5 * acceX * delta_t * delta_t
        #     Pinit_y = Pinit_y + Vinit_y * delta_t + 0.5 * acceY * delta_t * delta_t
        #     Pinit_z = Pinit_z + Vinit_z * delta_t + 0.5 * (acceZ) * delta_t * delta_t

        ### gravity direction problem, the original version of getting V and P after KF processed
        # Vinit_x = Vinit_x + paramAcce_x[0][0][0] * delta_t
        # Vinit_y = Vinit_y + paramAcce_y[0][0][0] * delta_t
        # Vinit_z = Vinit_z + (paramAcce_z[0][0][0]) * delta_t
        #
        # Pinit_x = Pinit_x + Vinit_x * delta_t + 0.5 * paramAcce_x[0][0][0] * delta_t * delta_t
        # Pinit_y = Pinit_y + Vinit_y * delta_t + 0.5 * paramAcce_y[0][0][0] * delta_t * delta_t
        # Pinit_z = Pinit_z + Vinit_z * delta_t + 0.5 * (paramAcce_z[0][0][0]) * delta_t * delta_t

        #### Integrate to V then KF
        # Adiff_x, Adiff_y, Adiff_z, Pinit_x, Pinit_y, Pinit_z, LastV_x, LastV_y, LastV_z, LastP_x, LastP_y, LastP_z\
        #     = filterVafterIntegrate(paramAcce_x[0][0][0], paramAcce_y[0][0][0], paramAcce_z[0][0][0],
        #                             LastV_x, LastV_y, LastV_z, LastP_x, LastP_y, LastP_z)
        #### Integrate to P then KF
        Adiff_x, Adiff_y, Adiff_z, Vdiff_x, Vdiff_y, Vdiff_z, LastV_x, LastV_y, LastV_z, LastP_x, LastP_y, LastP_z\
            = filterPafterIntegrate(paramAcce_x[0][0][0], paramAcce_y[0][0][0], paramAcce_z[0][0][0],
                                    LastV_x, LastV_y, LastV_z, LastP_x, LastP_y, LastP_z)



        # csv_writer.writerow([wx, wy, wz, qt[0][0], qt[1][0], qt[2][0], qt[3][0], paramAcce_x[0][0][0],
        #                       paramAcce_y[0][0][0], paramAcce_z[0][0][0], rawAcce_x, rawAcce_y, rawAcce_z, Pinit_x,
        #                       Pinit_y, Pinit_z, RawP_x, RawP_y, RawP_z])
        ### Write after integrate to V then KF
        # csv_writer.writerow([wx, wy, wz, qt[0][0], qt[1][0], qt[2][0], qt[3][0], Adiff_x,
        #                      Adiff_y, Adiff_z, rawAcce_x, rawAcce_y, rawAcce_z, Pinit_x, Pinit_y, Pinit_z, RawP_x,
        #                      RawP_y, RawP_z])
        ### Write after integrate to P then KF
        csv_writer.writerow([wx, wy, wz, qt[0][0], qt[1][0], qt[2][0], qt[3][0], Adiff_x, Adiff_y, Adiff_z,
                             rawAcce_x, rawAcce_y, rawAcce_z, paramAcce_x[0][0][0], paramAcce_y[0][0][0], paramAcce_z[0][0][0],
                             RawP_x, RawP_y, RawP_z])

        Q_before = qt
        # print(Q_before)

### return under integrate to V then KF && Acce
    # return [Pinit_x, Pinit_y, Pinit_z, RawP_x, RawP_y, RawP_z, initialize[4], initialize[5], initialize[6], getGroup[5][0],
    #         getGroup[5][1], getGroup[5][2]]
### return under integrate to P then KF
    return [paramAcce_x[0][0][0], paramAcce_y[0][0][0], paramAcce_z[0][0][0], RawP_x, RawP_y, RawP_z, initialize[4], initialize[5], initialize[6], getGroup[5][0],
            getGroup[5][1], getGroup[5][2]]